"""
眼疾识别
"""
import os
import random

import cv2
import numpy as np
import matplotlib.pyplot as plt
import paddle
from PIL import Image
from paddle.vision.models import resnet50
import paddle.nn.functional as F


# DATADIR = 'data/PALM-Training400/PALM-Training400'
# # 文件名以N开头的是正常眼底照片，以P开头的是病变眼底照片
# file1 = 'N0012.jpg'
# file2 = 'P0095.jpg'
#
# # 读取照片
# img1 = Image.open(os.path.join(DATADIR,file1))
# img1 = np.array(img1)
# img2 = Image.open(os.path.join(DATADIR,file2))
# img2 = np.array(img2)
#
# # 画出读取的照片
# plt.figure(figsize=(16,8))
# f = plt.subplot(121)
# f.set_title('Normal',fontsize=20)
# plt.imshow(img1)
# f = plt.subplot(122)
# f.set_title('PM',fontsize=20)
# plt.imshow(img2)
# plt.show()

def transform_img(img):
    # 将图片尺寸缩放到 224x224
    img = cv2.resize(img,(224,224))
    # 读入的图像数据格式是[H,W,C]
    # 使用转置操作将其变为[C,H,W]
    img = np.transpose(img,(2,0,1))
    img = img.astype('float32')

    img = img/255.
    img = img * 2.0 -1.0
    return img

def data_loader(datadir,batch_size=10,mode='train'):
    # 将datadir目录下的文件列出来，每条文件都要读入
    filenames = os.listdir(datadir)
    def reader():
        if mode == 'train':
            # 训练时随机打乱数据顺序
            random.shuffle(filenames)
        batch_imgs = []
        batch_labels = []
        for name in filenames:
            filepath = os.path.join(datadir,name)
            img = cv2.imread(filepath)
            img = transform_img(img)
            if name[0] == 'H' or name[0] == 'N':
                # H开头的文件名表示高度近似，N开头的文件名表示正常视力
                # 高度近视和正常视力的样本，都不是病理性的，属于负样本，标签为0
                label = 0
            elif name[0] == 'P':
                # P开头的是病理性近视，属于正样本，标签为1
                label = 1
            else:
                raise('Not excepted file name')
            # 没读取一个样本的数据，就将其放入数据列表中
            batch_imgs.append(img)
            batch_labels.append(label)
            if len(batch_imgs) == batch_size:
                # 当数据列表的长度等于batch_size的时候
                # 把这些数据当作一个mini_batch,并作为数据生成器的一个输出
                imgs_array = np.array(batch_imgs).astype('float32')
                labels_array = np.array(batch_labels).reshape(-1,1)
                yield imgs_array,labels_array
                batch_imgs = []
                batch_labels = []

        if len(batch_imgs) > 0:
            # 剩余样本数目不足一个batch_size的数据，一起打包成一个mini_batch
            imgs_array = np.array(batch_imgs).astype('float32')
            labels_array = np.array(labels_array).reshape(-1,1)
            yield imgs_array,labels_array

    return reader

def valid_data_loader(datadir,csvfile,batch_size=10,mode='valid'):
    # 训练集读取时通过文件名来确定样本标签，验证集则通过csvfile来读取每个图片对应的标签
    # 请查看解压后的验证集标签数据，观察csvfile文件里面所包含的内容
    # csvfile文件所包含的内容格式如下，每一行代表一个样本
    # 其中第一列是图片id,第二列是文件名，第三列是图片标签
    # 第四列和第五列是Fovea的坐标，与分类任务无关
    # ID,imgName,Label,Fovea_X,Fovea_Y
    # 1,V0001.jpg,0,1157.74,1019.87
    # 2,V0002.jpg,1,1285.82,1080.47
    filelists = open(csvfile).readlines()
    def reader():
        batch_imgs = []
        batch_labels = []
        for line in filelists[1:]:
            line = line.strip().split(',')
            # print(line)
            name = line[1]
            label = int(line[2])
            # 根据图片文件名加载图片，并对图像数据作预处理
            filepath = os.path.join(datadir,name)
            img = cv2.imread(filepath)
            img = transform_img(img)
            # 每读取一个样本的数据，就将其放入数据列表中
            batch_imgs.append(img)
            batch_labels.append(label)
            if len(batch_imgs) == batch_size:
                # 当数据列表的长度等于batch_size的时候
                # 把这些数据当作一个mini-batch,并作为数据生成器的一个输出
                imgs_array = np.array(batch_imgs).astype('float32')
                labels_array = np.array(batch_labels).reshape(-1,1)
                yield imgs_array,labels_array
                batch_imgs = []
                batch_labels = []

        if len(batch_imgs) > 0:
            # 剩余样本数目不足一个batch_size的价值，一起打包成一个mini_batch
            imgs_array = np.array(batch_imgs).astype('float32')
            labels_array = np.array(batch_labels).reshape(-1,1)
            yield imgs_array,labels_array
    return reader

paddle.seed(100)

# 查看数据形状
# DATADIR = 'data/PALM-Training400/PALM-Training400'
# train_loader = data_loader(DATADIR,batch_size=10,mode='train')
# data_reader = train_loader()
# data = next(data_reader)
# print(data[0].shape,data[1].shape)
#
# eval_loader = data_reader(DATADIR,batch_size=10,mode='eval')
# data_reader = eval_loader()
# data = next(data_reader)
# print(data[0].shape,data[1].shape)

model = resnet50()
loss_fn = F.cross_entropy

class Runner(object):
    def __init__(self,model,optimizer,loss_fn):
        self.model = model
        self.optimizer = optimizer
        self.loss_fn = loss_fn

        # 记录全局最优指标
        self.best_acc = 0

    # 定义训练过程
    def train_pm(self,train_datadir,val_datadir,**kwargs):
        print('start training ...')
        self.model.train()

        num_epochs = kwargs.get('num_epochs',0)
        csv_file = kwargs.get('csv_file',0)
        save_path = kwargs.get("save_path","output/")

        # 定义数据读取器，训练数据读取器
        train_loader = data_loader(train_datadir,batch_size=10,mode='train')

        for epoch in range(num_epochs):
            for batch_id, data in enumerate(train_loader()):
                x_data, y_data = data
                img = paddle.to_tensor(x_data)
                label = paddle.to_tensor(y_data)
                # 运行模型前向计算，得到预测值
                logits = model(img)
                avg_loss = self.loss_fn(logits,label)

                if batch_id % 20 == 0:
                    print("epoch: {}, batch_id: {}, loss is: {:.4f}".format(epoch,batch_id,float(avg_loss.numpy())))

                # 反向传播，更新权重，清除梯度
                avg_loss.backward()
                self.optimizer.step()
                self.optimizer.clear_grad()

            acc = self.evaluate_pm(val_datadir,csv_file)
            self.model.train()
            if acc > self.best_acc:
                self.save_model(save_path)
                self.best_acc = acc


    # 模型评估阶段，使用'paddle.no_grad()'控制不计算和存储梯度
    @paddle.no_grad()
    def evaluate_pm(self,val_datadir,csv_file):
        self.model.eval()
        accuracies = []
        losses = []
        # 验证数据读取器
        valid_loader = valid_data_loader(val_datadir,csv_file)

        for batch_id,data in enumerate(valid_loader()):
            x_data, y_data = data
            img = paddle.to_tensor(x_data)
            label = paddle.to_tensor(y_data)
            # 运行模型前向计算，得到预测值
            logits = self.model(img)
            # 多分类，使用softmax计算预测概率
            pred = F.softmax(logits)
            loss = self.loss_fn(pred,label)
            acc = paddle.metric.accuracy(pred,label)
            accuracies.append(acc.numpy())
            losses.append(loss.numpy())
        print("[validation] accuracy/loss: {:.4f}/{:.4f}".format(np.mean(accuracies),np.mean(losses)))
        return np.mean(accuracies)

    # 模型评估阶段，使用'paddle.no_grad()'控制不计算和存储梯度
    @paddle.no_grad()
    def predict_pm(self,x,**kwargs):
        # 将模型设置为评估模式
        self.model.eval()
        # 运行模型前向计算，得到预测值
        logits = self.model(x)
        return logits

    def save_model(self,save_path):
        paddle.save(self.model.state_dict(),save_path + 'palm.pdparams')
        paddle.save(self.optimizer.state_dict(),save_path + 'palm.pdopt')

    def load_model(self,model_path):
        model_state_dict = paddle.load(model_path)
        self.model.set_state_dict(model_state_dict)

# 开启0号GPU训练
use_gpu = True
paddle.device.set_device('gpu:0') if use_gpu else paddle.device.set_device('cpu')

# 定义优化器
# opt = paddle.optimizer.Momentum(learning_rate=0.001,momentum=0.9,parameters=model.parameters(),weight_decay=0.001)
opt = paddle.optimizer.SGD(learning_rate=0.001,parameters=model.parameters())

runner = Runner(model,opt,loss_fn)

# 数据集路径
DATADIR = 'data/PALM-Training400/PALM-Training400'
DATADIR2 = 'data/PALM-Validation400'
CSVFILE = 'data/PALM-Validation-GT/PM_Label_and_Fovea_Location.csv'

# 设置迭代轮数
EPOCH_NUM = 5
# 模型保存路径
PATH='output/'
if not os.path.exists(PATH):
    os.makedirs(PATH)

# 启动训练过程
runner.train_pm(DATADIR,DATADIR2,num_epochs=EPOCH_NUM,csv_file=CSVFILE,save_path=PATH)

# # 加载最优模型
# runner.load_model('output/palm.pdparams')
# # 模型评价
# score = runner.evaluate_pm(DATADIR2,CSVFILE)

# 模型预测
runner.load_model('output/palm.pdparams')

# 获取测试集中第一条数据
DATADIRv2 = 'data/PALM-Validation400'
filelists = open('data/PALM-Validation-GT/PM_Label_and_Fovea_Location.csv').readlines()
# 可以通过修改filelists列表的数字获取其他测试图片，可取值1-400
line = filelists[1].strip().split(',')
name,label = line[1],int(line[2])
print('name: {}, label: {}'.format(name,label))

# print(os.path.join(DATADIRv2,name))
# 读取测试图片
img = cv2.imread(os.path.join(DATADIRv2,name))
# 测试图片预处理
trans_img = transform_img(img)
unsqueeze_img = paddle.unsqueeze(paddle.to_tensor(trans_img),axis=0)

# 模型预测
logits = runner.predict_pm(unsqueeze_img)
result = F.softmax(logits)
pred_class = paddle.argmax(result).numpy()

# 输出真实类别和预测类别
print("The true category is {} and the predicted category is {}".format(label,pred_class))

# 图片可视化
show_img = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
plt.imshow(show_img)
plt.show()